# encoding: utf-8
# @File  : mongo_manager.py
# @Author: shaoyun
# @Date  :  2025/05/11
import asyncio
from typing import Dict, List, Optional, Any

from motor.motor_asyncio import AsyncIOMotorClient
from pymongo import MongoClient
from pymongo.errors import PyMongoError

from app.core.config import settings
from loguru import logger


class MongoDB:
    _async_client: Optional[AsyncIOMotorClient] = None
    _sync_client: Optional[MongoClient] = None
    _lock = asyncio.Lock()

    @classmethod
    async def connect_async(cls):
        """连接到MongoDB（异步）"""
        async with cls._lock:
            if cls._async_client is None:
                try:
                    cls._async_client = AsyncIOMotorClient(
                        settings.MONGODB_URL,
                        maxPoolSize=10,
                        minPoolSize=1
                    )
                    # 测试连接
                    await cls._async_client.admin.command('ping')
                    logger.info("Connected to MongoDB (async)")
                except PyMongoError as e:
                    logger.error(f"Failed to connect to MongoDB (async): {e}")
                    raise

    @classmethod
    def connect_sync(cls):
        """连接到MongoDB（同步）"""
        if cls._sync_client is None:
            try:
                cls._sync_client = MongoClient(
                    settings.MONGODB_URL,
                    maxPoolSize=10,
                    minPoolSize=1
                )
                # 测试连接
                cls._sync_client.admin.command('ping')
                logger.info("Connected to MongoDB (sync)")
            except PyMongoError as e:
                logger.error(f"Failed to connect to MongoDB (sync): {e}")
                raise

    @classmethod
    async def close_async(cls):
        """关闭MongoDB连接（异步）"""
        if cls._async_client:
            cls._async_client.close()
            cls._async_client = None
            logger.info("Closed MongoDB connection (async)")

    @classmethod
    def close_sync(cls):
        """关闭MongoDB连接（同步）"""
        if cls._sync_client:
            cls._sync_client.close()
            cls._sync_client = None
            logger.info("Closed MongoDB connection (sync)")

    @classmethod
    async def get_async_client(cls) -> AsyncIOMotorClient:
        """获取异步客户端，如果未连接则自动连接"""
        if cls._async_client is None:
            await cls.connect_async()
        return cls._async_client

    @classmethod
    def get_sync_client(cls) -> MongoClient:
        """获取同步客户端，如果未连接则自动连接"""
        if cls._sync_client is None:
            cls.connect_sync()
        return cls._sync_client

    @classmethod
    async def get_database(cls, async_mode=True):
        """获取数据库对象"""
        if async_mode:
            client = await cls.get_async_client()
            return client[settings.MONGODB_DB]
        else:
            client = cls.get_sync_client()
            return client[settings.MONGODB_DB]

    @classmethod
    async def insert_one(cls, collection: str, document: Dict) -> str:
        """插入单个文档并返回ID"""
        try:
            db = await cls.get_database()
            result = await db[collection].insert_one(document)
            return str(result.inserted_id)
        except PyMongoError as e:
            logger.error(f"Error inserting document into {collection}: {e}")
            raise

    @classmethod
    async def find_one(cls, collection: str, query: Dict) -> Optional[Dict]:
        """查找单个文档"""
        try:
            db = await cls.get_database()
            result = await db[collection].find_one(query)
            return result
        except PyMongoError as e:
            logger.error(f"Error finding document in {collection}: {e}")
            raise

    @classmethod
    async def find_many(cls, collection: str, query: Dict, skip: int = 0, limit: int = 100) -> List[Dict]:
        """查找多个文档"""
        try:
            db = await cls.get_database()
            cursor = db[collection].find(query).skip(skip).limit(limit)
            return await cursor.to_list(length=limit)
        except PyMongoError as e:
            logger.error(f"Error finding documents in {collection}: {e}")
            raise

    @classmethod
    async def update_one(cls, collection: str, query: Dict, update: Dict) -> int:
        """更新单个文档"""
        try:
            db = await cls.get_database()
            result = await db[collection].update_one(query, {"$set": update})
            return result.modified_count
        except PyMongoError as e:
            logger.error(f"Error updating document in {collection}: {e}")
            raise

    @classmethod
    async def delete_one(cls, collection: str, query: Dict) -> int:
        """删除单个文档"""
        try:
            db = await cls.get_database()
            result = await db[collection].delete_one(query)
            return result.deleted_count
        except PyMongoError as e:
            logger.error(f"Error deleting document from {collection}: {e}")
            raise

    @classmethod
    async def count_documents(cls, collection: str, query: Dict) -> int:
        """统计文档数量"""
        try:
            db = await cls.get_database()
            return await db[collection].count_documents(query)
        except PyMongoError as e:
            logger.error(f"Error counting documents in {collection}: {e}")
            raise

    @classmethod
    async def aggregate(cls, collection: str, pipeline: List[Dict]) -> List[Dict]:
        """聚合查询"""
        try:
            db = await cls.get_database()
            cursor = db[collection].aggregate(pipeline)
            return await cursor.to_list(length=None)
        except PyMongoError as e:
            logger.error(f"Error aggregating in {collection}: {e}")
            raise